Actionable comments posted: 2

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Around line 59-60: The current enablement enables HybridModel defaults but
pattern updates and metric accounting are still gated by isinstance(...,
MambaModel), causing plain HybridModel instances to skip hybrid-pattern logic;
update the checks in the functions that perform hybrid pattern updates and
candidate metric accounting (the places currently using isinstance(obj,
MambaModel)) to detect HybridModel instead (e.g., isinstance(obj, HybridModel))
so plain HybridModel instances get the same pattern handling and metric updates,
and if any MambaModel-specific behavior is required keep a secondary
isinstance(obj, MambaModel) branch for those special cases.

In `@modelopt/torch/utils/dataset_utils.py`:
- Around line 492-526: The code can silently produce zero-length output when
token_stream is too short: after computing n_chunks from token_stream, check for
the edge cases and fail fast or warn; specifically, in the packing branch after
computing n_chunks (and before building input_ids/batch_encoded) add logic that
raises a clear exception if n_chunks == 0 (mentioning tokenizer.encode,
token_stream, total_chunks, and max_sample_length) and log/warn when n_chunks <
total_chunks to inform the caller they received fewer chunks than requested;
ensure the exception/warning uses the existing logging mechanism or raises a
ValueError so callers cannot proceed with empty tensors.

Claude review summary

Findings: CRITICAL: 1, IMPORTANT: 1, SUGGESTION: 3.

The shared calibration loop is a clean consolidation and the experimental data backs the design. One blocker, one back-compat side effect, and a few smaller notes — all in the new megatron_calibration.py.

Highest-impact

• CRITICAL Algorithm — padding-direction assumption (megatron_calibration.py:131-144). The per-row trim slices ids[b, :real_len] and the no-padding branch overwrites ids[:, -1], both of which assume right-padded sequences. But get_dataset_dataloader (the function this loop calls) explicitly recommends padding_side="left" and warns when it isn't (dataset_utils.py:593-596). With left-padding, the trim selects padding tokens instead of real tokens — re-introducing the exact contamination the PR is trying to eliminate. Your Qwen3-8B numbers move in the right direction because that tokenizer happens to default to right-padding, but a user who follows the dataloader's documented recommendation silently regresses. Capture padding_side once and branch the slice (ids[b, -real_len:] for left-padding); also worth a regression test row whose attention_mask has zeros to lock the behavior in.

• IMPORTANT Compatibility — caller-tokenizer mutation (megatron_calibration.py:87-88). Setting tokenizer.pad_token = tokenizer.eos_token mutates the caller's tokenizer object before the dataloader's deepcopy. The unwrapped tokenizer in prune_minitron.py is reused downstream (MMLU eval, export). Tokenizers whose pad_token was None will silently start emitting EOS as their pad token in every later code path that touches them. Local-copy the tokenizer before the mutation, or set the field on the dataloader-internal deepcopy only.

SUGGESTIONs (non-blocking)

• Docstring claim that the loop matches GPTSFTDataset(add_eos=True) is slightly off — that flag appends EOS, the code overwrites the last real token. Reword or actually append when under seq_length.
• bool((mask == 0).any().item()) and mask[b].sum().item() cause a CPU-GPU sync per batch / per row. Precompute real-lengths once on device into a CPU list at builder time to drop ~2*num_batches syncs.
• The CP-rank slice inside _forward_loop is dead code in the tested configs (CP=1) and would actually break under CP>1 because different CP ranks would compute different per-row real_len and call megatron_prefill (a collective) with shape-mismatched inputs. Either run the trim before the CP slice, or drop the CP-shard call until CP>1 is validated.

Backward compatibility

• Public API change: removing get_hf_mbridge_calibration_loop from modelopt.torch.utils.plugins.mbridge is a breaking change for any out-of-tree caller. Acceptable since the in-tree call sites are migrated and there's no schema/checkpoint impact, but flagging it for the maintainer's awareness.
• examples/megatron_bridge/prune_minitron.py renames --calib_mbs / --calib_gbs to --calib_batch_size. Existing user scripts/CI that invoke this example with the old flags will fail at parse time. README is updated in the same PR — fine.

Risk

Low-to-moderate. The CRITICAL is a correctness contradiction with the dataloader the loop relies on; once that's reconciled and the tokenizer mutation is confined to a local copy, this is a clear quality win over the previous bespoke loops.

Claude review summary

Findings: CRITICAL: 0, IMPORTANT: 2, SUGGESTION: 2

The refactor is well-motivated: the unified loop is simpler than the old M-Bridge evaluate_and_print_results pipeline, the per-row trim + EOS-at-end matches GPTSFTDataset(add_eos=True), and the sort-by-length / per-batch padding-fast-path is a clean way to get batched throughput without polluting calibration stats with padding. The MMLU numbers in the PR description back this up.

Most impactful findings

1. Public API removal without deprecation — get_hf_mbridge_calibration_loop is dropped from modelopt.torch.utils.plugins.mbridge.__all__ outright. Out-of-tree callers that pinned a previous example version will hit ImportError on upgrade with no migration hint. A one-release deprecation shim is cheap.

2. inference_batch_size defaulting to calib_batch_size silently changes prune_target_memory_mb math. Previously --calib_mbs was hard-pinned to 1; now --calib_batch_size is unrestricted and the PR description recommends 16 for throughput. A user who follows that recommendation without also setting --inference_batch_size gets a 16× larger KV-cache term in the memory constraint — i.e., a different prune target — silently. Either default inference_batch_size to 1 unconditionally, or warn loudly when this divergence is implied.

Non-blocking

3. CP > 1 is not really supported by the get_batch_on_this_cp_rank + megatron_prefill combination (megatron_prefill rebuilds attention_mask / position_ids from the locally-sliced length). Fine while the validated configs are CP=1, worth either an assert or removing the CP wrapper to avoid future foot-gun.
4. Minor docstring nit: "no logits compute" overstates what skip_return_logits=True does.

Risk assessment

Algorithmic changes are validated by the MMLU sweeps in the PR description and the per-row trim / EOS / no-padding-in-calibration logic looks right. The risk surface is the user-facing surface: removed public symbol and the silent-default-shift on the prune memory target. Both are easy to address.

Claude re-review

Findings (this pass): CRITICAL: 0, IMPORTANT: 2 (both carried over from prior review, still unresolved in code), SUGGESTION: 0.

Re-verification of prior findings

The single new commit (Address Claude review: CP-assert + clarify prefill docstring) addressed the two SUGGESTION items from the 2026-05-20 review:

• ✅ CP > 1 footgun — runtime RuntimeError now raised inside _forward_loop when context_parallel_world_size != 1.
• ✅ Docstring overstatement — module docstring now accurately describes skip_return_logits=True as "skips returning logits / loss compute … the LM head still runs and activation hooks still fire."

The two IMPORTANT items are not addressed:

1. [IMPORTANT Compatibility] get_hf_mbridge_calibration_loop is removed from modelopt/torch/utils/plugins/mbridge.py:__all__ outright with the implementation deleted. Out-of-tree callers pinned to a previous release will hit ImportError on upgrade with no migration hint. A one-release deprecation shim (re-export plus DeprecationWarning pointing to get_megatron_calibration_forward_loop) would be cheap insurance.

2. [IMPORTANT Compatibility] examples/megatron_bridge/prune_minitron.py:155-164 defaults --inference_batch_size to --calib_batch_size. Previously --calib_mbs was hard-pinned to choices=[1], so inference_batch_size always fell back to 1 — users could safely bump --calib_gbs for throughput without affecting the prune memory target. Now, the PR description recommends --calib_batch_size 16 (and the experimental table in the description used calib_bs 16 for the production runs), but anyone following that recommendation without also setting --inference_batch_size 1 will silently get a 16× larger KV-cache term in --prune_target_memory_mb — a different prune target than they intended. Default inference_batch_size=1 unconditionally, or warn when divergence is implied.

Algorithm / mode-state re-trace

End-to-end trace of the new get_megatron_calibration_forward_loop against current state:

• Padding direction — boolean-mask gather (ids[b][mask[b].bool()]) correctly handles both left- and right-padded tokenizers; the prior padding-direction concern is fully addressed.
• Tokenizer mutation — local copy.deepcopy(tokenizer) before tokenizer.pad_token = tokenizer.eos_token correctly isolates the mutation from the caller's tokenizer. get_dataset_dataloader does its own deepcopy, so the caller is double-protected.
• CPU-GPU sync — lengths_cpu precomputed once on CPU; no per-batch sync inside the forward hot loop. ✓
• Sort-by-length / per-batch padding-fast-path — for batches with no padding, batched forward; for any batch containing padding, falls through to per-row forward. Calibration statistics (amax / channel-importance) are order-invariant aggregates, so the re-ordering is bit-identical to un-sorted. ✓
• PP coordination — all PP ranks see the same dataloader output (deterministic get_dataset_samples + shuffle=False) and call megatron_prefill with matching shapes per iteration. ✓
• EOS-at-row-end semantics — overwrites the row's last real token with EOS for both the per-row and batched paths. Docstring is honest about the under-cap-row content-token loss trade-off. ✓
• Loop invariants — partial-tail batches (n % batch_size != 0) handled correctly via for b in range(ids.shape[0]). Zero-real-token rows skipped via row.shape[1] < 1 continue.

Backward compatibility

• CLI rename --calib_mbs/--calib_gbs → --calib_batch_size (item 2 above).
• Public API removal of get_hf_mbridge_calibration_loop (item 1 above).
• No modelopt_state schema changes.

Risk

Low-to-moderate. The algorithm itself is sound and matches the experimental data in the PR description. Risk surface is the user-facing surface: a public-symbol removal with no shim, and a CLI default coupling that quietly changes the prune target when the user follows the PR's own throughput recommendation. Both are easy to address.